import shutil
import tkinter as tk
from  tkinter import filedialog, ttk, messagebox, font

from PIL import Image, ImageTk

import os
import cv2

from pathlib import Path


class App(tk.Frame):
    def __init__(self, master):
        super().__init__(master)
        self.master.minsize(640, 540)
        self.master.maxsize(self.winfo_screenwidth(), self.winfo_screenheight())
        self.place(relx=0, rely=0, relwidth=1, relheight=1)

        self.contents_open_path = tk.StringVar()
        self.contents_open_video = tk.StringVar()
        self.contents_save_path = tk.StringVar()

        self.contents_video_num = tk.StringVar()
        self.contents_video_num.set('共检测到   个视频')

        # self.res_width = tk.StringVar()
        # self.res_height = tk.StringVar()

        self._content_pasue = tk.BooleanVar()
        self._content_pasue.set(False)
        self._content_stop = tk.BooleanVar()
        self._content_stop.set(False)

        lf_path = tk.LabelFrame(self, text='路径设置')
        lf_path.place(relx=0.01, rely=0.01, relwidth=0.98, relheight=0.2)

        self.butt_open_path = tk.Button(lf_path)
        self.butt_open_path['text'] = "打开文件夹"
        self.butt_open_path['command'] = self._open_dir
        self.butt_open_path.place(relx=0.01, rely=0.01, relwidth=0.18, relheight=0.28)

        label_op = tk.Label(lf_path, bg='white')
        label_op["textvariable"] = self.contents_open_path
        label_op.place(relx=0.2, relwidth=0.6, rely=0.01, relheight=0.28)

        label_n = tk.Label(lf_path)
        label_n["textvariable"] = self.contents_video_num
        label_n.place(relx=0.8, relwidth=0.2, rely=0.01, relheight=0.28)

        self.butt_open_video = tk.Button(lf_path)
        self.butt_open_video['text'] = "打开视频"
        self.butt_open_video['command'] = self._open_video
        self.butt_open_video.place(relx=0.01, relwidth=0.18, rely=0.34, relheight=0.28)

        label_ov = tk.Label(lf_path, bg='white')
        label_ov["textvariable"] = self.contents_open_video
        label_ov.place(relx=0.2, relwidth=0.6, rely=0.34, relheight=0.28)

        self.butt_save_path = tk.Button(lf_path)
        self.butt_save_path['text'] = "保存到"
        self.butt_save_path['command'] = self._save_dir
        self.butt_save_path.place(relx=0.01, relwidth=0.18, rely=0.67, relheight=0.28)

        label_sp = tk.Label(lf_path, bg='white')
        label_sp["textvariable"] = self.contents_save_path
        label_sp.place(relx=0.2, relwidth=0.6, rely=0.67, relheight=0.28)

        label_step = tk.Label(lf_path, text='跳帧设置')
        label_step.place(relx=0.8, relwidth=0.1, rely=0.34, relheight=0.28)

        self.cmb_step = ttk.Combobox(lf_path)
        self.cmb_step['value'] = (1, 5, 15, 30, 60)
        self.cmb_step.place(relx=0.9, relwidth=0.09, rely=0.34, relheight=0.28)
        self.cmb_step.current(2)

        self.butt_open_save_path = tk.Button(lf_path)
        self.butt_open_save_path['text'] = "打开保存文件夹"
        self.butt_open_save_path['command'] = self._open_save_dir
        self.butt_open_save_path.place(relx=0.81, relwidth=0.18, rely=0.67, relheight=0.28)

        lf_convert = tk.LabelFrame(self, text='转换')
        lf_convert.place(relx=0.01, rely=0.23, relwidth=0.3, relheight=0.76)

        font_button = font.Font(size=36, weight='bold')
        self.butt_start = tk.Button(lf_convert, text="开始", font=font_button)
        self.butt_start['state'] = 'disabled'
        self.butt_start['command'] = self._start_convert
        self.butt_start.place(relx=0.1, rely=0.1, relwidth=0.8, relheight=0.2)

        self.butt_pause = tk.Button(lf_convert, text="暂停", command=self._pause_convert, font=font_button)
        self.butt_pause['state'] = 'disabled'
        self.butt_pause.place(relx=0.1, rely=0.4, relwidth=0.8, relheight=0.2)

        self.butt_stop = tk.Button(lf_convert, text="重置", command=self._stop_convert, font=font_button)
        self.butt_stop.place(relx=0.1, rely=0.7, relwidth=0.8, relheight=0.2)
        self.butt_stop['state'] = 'disabled'

        lf_res = tk.LabelFrame(self, text='保存分辨率')
        lf_res.place(relx=0.32, relwidth=0.67, rely=0.23, relheight=0.1)

        self.cmb_res_choices = ttk.Combobox(lf_res)
        self.cmb_res_choices['value'] = (
            '   origin', 
            '   720p', 
            '   user_defined')
        self.cmb_res_choices.state(['readonly'])
        self.cmb_res_choices.place(relx=0.01, relwidth=0.29, rely=0.05, relheight=0.8)
        self.cmb_res_choices.current(0)
        # self.cmb_res_choices.bind('   720p', self._choose_res)
        self.cmb_res_choices.bind('<<ComboboxSelected>>', self._choose_res)

        self.var_res_width = tk.StringVar()
        self.var_res_height = tk.StringVar()

        self.res_width = ttk.Entry(lf_res, textvariable=self.var_res_width)
        self.res_width.place(relx=0.47, relwidth=0.2, rely=0.05, relheight=0.8)
        self.res_width['state'] = 'disabled'
        self.res_height = ttk.Entry(lf_res, textvariable=self.var_res_height)
        self.res_height.place(relx=0.79, relwidth=0.2, rely=0.05, relheight=0.8)
        self.res_height['state'] = 'disabled'

        lb_res_w = tk.Label(lf_res, text='width:')
        lb_res_w.place(relx=0.35, relwidth=0.1, rely=0.05, relheight=0.8)

        lb_res_h = tk.Label(lf_res, text='height:')
        lb_res_h.place(relx=0.67, relwidth=0.11, rely=0.05, relheight=0.8)

        self.lf_prog = tk.LabelFrame(self, text='进度条')
        self.lf_prog.place(relx=0.32, rely=0.34, relwidth=0.67, relheight=0.65)

        self.label_prog1 = tk.Label(self.lf_prog, text='总进度条：已完成 ___ | ___')
        self.label_prog1.place(relx=0.01, relwidth=0.98, rely=0, relheight=0.05)

        self.prog_1 = ttk.Progressbar(self.lf_prog, maximum=500)
        self.prog_1.place(relx=0.01, rely=0.05, relwidth=0.98, relheight=0.08)

        self.label_prog2 = tk.Label(self.lf_prog, text='')
        self.label_prog2.place(relx=0.01, relwidth=0.98, rely=0.13, relheight=0.05)

        self.prog_2 = ttk.Progressbar(self.lf_prog)
        self.prog_2.place(relx=0.01, rely=0.18, relwidth=0.98, relheight=0.08)

        self.lb_image = tk.Label(self.lf_prog, image=None, bg='light yellow')
        self.lb_image.place(relx=0.0, rely=0.3, relwidth=1, relheight=0.7)
    
    def _choose_res(self, event=None):
        choice = self.cmb_res_choices.get()
        if choice in ['   origin']:
            self.var_res_width.set('')
            self.var_res_height.set('')
            self.res_width['state'] = 'readonly'
            self.res_height['state'] = 'readonly'
        if choice in ['   720p']:
            self.var_res_width.set('1280')
            self.var_res_height.set('720')
            self.res_width['state'] = 'readonly'
            self.res_height['state'] = 'readonly'
        if choice in ['   user_defined']:
            self.res_width['state'] = 'normal'
            self.res_height['state'] = 'normal'

    
    def _reset(self):
        self.prog_1['value'] = 0
        self.prog_2['value'] = 0
        self.butt_start['text'] = '开始'
        self.butt_start['state'] = 'normal'
        self.butt_pause['state'] = 'disabled'
        self.butt_stop['state'] = 'disabled'
        self.label_prog1['text'] = '总进度条：已完成 ___ | ___'
        self.label_prog2['text'] = ''
        self.butt_open_path['state'] = 'normal'
        self.butt_open_video['state'] = 'normal'
        self.butt_save_path['state'] = 'normal'
        self.cmb_step['state'] = 'normal'
        self._choose_res()
        self.lb_image.imgtk = None
    
    def _start_convert(self):
        self.butt_start['text'] = '运行中'
        self.butt_start['state'] = 'disabled'
        self.butt_pause['state'] = 'normal'
        self.butt_stop['state'] = 'normal'

        self.butt_open_path['state'] = 'disabled'
        self.butt_open_video['state'] = 'disabled'
        self.butt_save_path['state'] = 'disabled'
        self.cmb_step['state'] = 'disabled'

        self.res_width['state'] = 'readonly'
        self.res_height['state'] = 'readonly'

        num_videos = len(self.valid_videos)
        step = 500 / num_videos
        self.prog_1['value'] = 0

        skip = False

        res_choice = self.cmb_res_choices.get()
        
        for i, path in enumerate(self.valid_videos):

            self.label_prog1['text'] = '总进度条：已完成 {:3d} | {:<3d}'.format(i, num_videos)
            
            try:
                # reader = FFMPEG_VideoReader(str(path))
                cap = cv2.VideoCapture(str(path))
                width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                nframes = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

                if res_choice in ['   origin']:
                    self.var_res_width.set(str(width))
                    self.var_res_height.set(str(height))
                    self.res_width['state'] = 'readonly'
                    self.res_height['state'] = 'readonly'

                valid_nframes = True
                if nframes <= 0:
                    valid_nframes = False
                    nframes = 300 * 60 * 25

                # img_ratio = reader.size[0] / reader.size[1]
                img_ratio = width / height
                frame_step = int(self.cmb_step.get())
                # self.prog_2['maximum'] = reader.nframes // frame_step
                self.prog_2['maximum'] = nframes if valid_nframes else 1
                self.prog_2['value'] = 0
                self.update()

                if self.open_path.is_dir():
                    save_dir = (Path(self.contents_save_path.get()) / path.relative_to(self.open_path)).with_suffix('')
                else:
                    save_dir = (Path(self.contents_save_path.get()) / path.relative_to(self.open_path))
                save_dir.mkdir(exist_ok=True, parents=True)

                self.label_prog2['text'] = str(path.relative_to(self.open_path))
                # for ind in range(reader.nframes):
                for ind in range(nframes):

                    if self._content_pasue.get():
                        self._content_pasue.set(False)
                        messagebox.showinfo('暂停','已暂停')
                        self.butt_pause['state'] = 'normal'
                    if self._content_stop.get():
                        self._content_stop.set(False)
                        ret = messagebox.askyesno('重置', '确定要重置吗？重置将删除已生成的图片及文件夹。', icon=messagebox.WARNING)
                        if ret:
                            self._reset()
                            for _path in self.valid_videos:
                                _save_dir = (Path(self.contents_save_path.get()) / _path.relative_to(self.open_path)).with_suffix('')
                                if _save_dir.exists():
                                    shutil.rmtree(str(_save_dir))
                            return None
                        else:
                            pass
                    
                    # frame = reader.read_frame()
                    valid, frame = cap.read()
                    if not valid:
                        break
                    
                    if valid_nframes:
                        self.prog_2['value'] += 1
                    else:
                        self.prog_2['value'] += 1 if ind % 2 == 0 else -1
                    self.update()

                    frame = frame[..., ::-1]
                    if ind % frame_step == 0:
                        # frame = reader.read_frame()
                        # reader.skip_frames(frame_step - 1)
                        win_size = (self.lb_image.winfo_width(), self.lb_image.winfo_height())
                        win_ratio = win_size[0] / win_size[1]
                        show_img = Image.fromarray(frame)
                        if win_ratio >= img_ratio:
                            img_h = win_size[1]
                            img_w = int(img_h * img_ratio)
                        else:
                            img_w = win_size[0]
                            img_h = int(img_w / img_ratio)
                        
                        show_img = show_img.resize((img_w, img_h))
                        imgtk = ImageTk.PhotoImage(show_img)
                        self.lb_image.imgtk = imgtk
                        self.lb_image.configure(image=imgtk)

                        # plt.imsave(save_dir / '{:05d}.jpg'.format(ind), frame)
                        ori_img = Image.fromarray(frame)
                        if res_choice in ['   720p', '   user_defined']:
                            ori_img = ori_img.resize((int(self.var_res_width.get()), int(self.var_res_height.get())))
                        ori_img.save(str(save_dir / '{:05d}.jpg'.format(ind)), quality=95)

            except Exception as e:
                if not skip:
                    ret = messagebox.askyesno(
                        '错误', f'文件：{str(path)}\n发生错误：{e}\n跳过并继续吗？', icon='error'
                    )
                    if ret:
                        continue
                    else:
                        self._reset()
                        break
                    ret = messagebox.askyesno(
                        '提示', f'下次跳过有错误文件，不再提示吗？', icon='info'
                    )
                    if ret:
                        skip = True
            
            self.label_prog1['text'] = '总进度条：已完成 {:3d} | {:<3d}'.format(i + 1, num_videos)
            self.prog_1['value'] += step
            self.update()
        
        # self.lb_image.imgtk = None
        self.butt_start['text'] = '已完成'
        self.butt_start['state'] = 'disabled'
        self.butt_pause['state'] = 'disabled'
        self.butt_stop['state'] = 'normal'
        self.butt_open_path['state'] = 'normal'
        self.butt_open_video['state'] = 'normal'
        self.butt_save_path['state'] = 'normal'
        self._choose_res()
    
    def _pause_convert(self):
        self.butt_pause['state'] = 'disabled'
        self._content_pasue.set(True)
    
    def _stop_convert(self):
        self._content_stop.set(True)
        if self.butt_start['text'] == '已完成':
            if messagebox.askyesno('重置','确定要重置吗？重置将删除已生成的图片及文件夹。', icon=messagebox.WARNING):
                self._reset()
                self._content_stop.set(False)
                for path in self.valid_videos:
                    save_dir = (Path(self.contents_save_path.get()) / path.relative_to(self.open_path)).with_suffix('')
                    shutil.rmtree(str(save_dir))
    
    def _open_save_dir(self):
        os.system('explorer {}'.format(Path(self.contents_save_path.get())))
    
    def _open_dir(self):
        self.open_path = Path(filedialog.askdirectory())
        self.contents_open_path.set(str(self.open_path))
        self.contents_open_video.set('')
        self.contents_save_path.set(str(self.open_path))
        self.all_videos = []
        for vfmt in ['mp4', 'avi']:
            self.all_videos += sorted(self.open_path.glob('**/*.{}'.format(vfmt)))
        existed_videos = []
        for path in self.all_videos:
            save_dir = path.with_suffix('')
            if save_dir.exists():
                existed_videos.append(path)
        
        self.valid_videos = self.all_videos.copy()
        if len(existed_videos) > 0:
            ret = messagebox.askyesnocancel(
                '重复保存警告', 
                f'已发现{len(existed_videos)}个与视频文件名同名的文件夹\n是否清空原文件夹？\n是：清空\n否：覆盖\n取消：跳过该文件\n还可以选择 否 或 取消 后重新设置保存路径'
            )
            if ret is True:
                for path in existed_videos:
                    save_dir = path.with_suffix('')
                    shutil.rmtree(str(save_dir))
            elif ret is None:
                self.valid_videos = [path for path in self.all_videos if path not in existed_videos]

        self.contents_video_num.set('共检测到{:3d}个视频'.format(len(self.all_videos)))
        
        self._reset()
        self.butt_start['state'] = 'normal' if len(self.valid_videos) else 'disabled'
    
    def _open_video(self):
        self.open_path = Path(filedialog.askopenfilename(filetypes=[('All video files', '.avi .mp4')]))
        self.contents_open_video.set(str(self.open_path))
        self.contents_open_path.set('')
        self.contents_save_path.set(str(self.open_path.with_suffix('')))
        self.contents_video_num.set('共检测到   个视频')
        
        self.all_videos = [self.open_path]
        self.valid_videos = self.all_videos.copy()
        
        save_dir = self.open_path.with_suffix('')
        if save_dir.exists():
            ret = messagebox.askyesnocancel(
                '重复保存警告', 
                f'已发现{str(self.open_path.name)}的同名文件夹\n是否清空原文件夹？\n是：清空\n否：覆盖\n取消：跳过该文件\n还可以选择 否 或 取消 后重新设置保存路径'
            )
            if ret is True:
                shutil.rmtree(str(save_dir))
            elif ret is None:
                self.valid_videos = []

        self._reset()
        self.butt_start['state'] = 'normal' if len(self.valid_videos) else 'disabled'

    def _save_dir(self):
        self.save_path = Path(filedialog.askdirectory())
        self.contents_save_path.set(str(self.save_path))

        existed_videos = []
        for path in self.all_videos:
            save_dir = (Path(self.contents_save_path.get()) / path.relative_to(self.contents_open_path.get())).with_suffix('')
            if save_dir.exists():
                existed_videos.append(path)
        
        self.valid_videos = self.all_videos.copy()
        if len(existed_videos) > 0:
            ret = messagebox.askyesnocancel(
                '重复保存警告', 
                f'已发现{len(existed_videos)}个与视频文件名同名的文件夹\n是否清空原文件夹？\n是：清空\n否：覆盖\n取消：跳过该文件\n还可以选择 否 或 取消 后重新设置保存路径'
            )
            if ret is True:
                for path in existed_videos:
                    save_dir = (Path(self.contents_save_path.get()) / path.relative_to(self.contents_open_path.get())).with_suffix('')
                    shutil.rmtree(str(save_dir))
            elif ret is None:
                self.valid_videos = [path for path in self.all_videos if path not in existed_videos]
        
        self.butt_start['state'] = 'normal' if len(self.valid_videos) else 'disabled'



if __name__ == '__main__':
    master = tk.Tk()
    master.iconbitmap('mv.ico')
    myapp = App(master)
    myapp.master.title("VIDEO2JPG v0.4 by Alpha@MOTOVIS")
    myapp.mainloop()